
from flask import Flask, jsonify, request, Blueprint
from flask_swagger_ui import get_swaggerui_blueprint
from typing import Union, Dict, Any, Callable, Optional, List
import pandas as pd
from nebula_data_engine.nebula_data_engine import NebulaDataEngine

class NebulaAPI:
    def __init__(
        self,
        data_source: Union[str, pd.DataFrame, Dict[str, pd.DataFrame], List[Dict[str, Any]], "NebulaDataEngine"],
        enable_docs: bool = True,
        enable_crud: bool = True,
        crud_class=None,
        auto_register: bool = True,  # 新增：是否自动注册路由
    ):
        """
        初始化 NebulaAPI。
        :param data_source: 数据源，可以是文件路径、MongoDB 连接字符串、字典、DataFrame、列表或 NebulaDataEngine 实例。
        :param enable_docs: 是否启用 Swagger 文档。
        :param enable_crud: 是否启用自动生成增删改查接口。
        :param crud_class: 自定义的 CRUD 类，用于覆盖默认的增删改查逻辑。
        :param auto_register: 是否自动注册路由。
        """
        self.app = Flask(__name__)
        self.enable_docs = enable_docs
        self.enable_crud = enable_crud
        self.crud_class = crud_class
        self.auto_register = auto_register  # 新增：是否自动注册路由
        self._swagger_paths = {}  # 用于存储 Swagger 文档的路由信息
        self._swagger_tags = []  # 用于存储 Swagger 文档的分组信息

        # 初始化 NebulaDataEngine
        if isinstance(data_source, NebulaDataEngine):
            self.data_engine = data_source
        else:
            self.data_engine = NebulaDataEngine(data_source)

        # 如果启用文档，配置 Swagger UI
        if self.enable_docs:
            self.setup_swagger()

        # 如果启用自动注册，注册 API 路由
        if self.auto_register:
            self.register_api_endpoints()

        self.bp = Blueprint('api', __name__, url_prefix='/api')

    def get_blueprint(self):
        return self.bp

    def get_flask_app(self):
        """获取 Flask 应用实例"""
        return self.app

    def setup_swagger(self):
        """配置 Swagger UI"""
        SWAGGER_URL = '/api/docs'  # Swagger UI 的访问路径
        API_URL = '/api/swagger.json'  # API 文档的 JSON 文件路径

        # 创建 Swagger UI 蓝图
        swaggerui_blueprint = get_swaggerui_blueprint(
            SWAGGER_URL,
            API_URL,
            config={
                'app_name': "NebulaAPI Documentation"
            }
        )

        # 注册 Swagger UI 蓝图
        self.app.register_blueprint(swaggerui_blueprint, url_prefix=SWAGGER_URL)

        # 定义 Swagger JSON 文件的路由
        @self.app.route(API_URL)
        def swagger_json():
            return jsonify({
                "openapi": "3.0.0",
                "info": {
                    "title": "NebulaAPI",
                    "version": "1.0.0",
                    "description": "自动生成的 NebulaAPI 文档"
                },
                "tags": self._swagger_tags,  # 添加分组信息
                "paths": self._swagger_paths
            })

    def generate_swagger_paths(self):
        """生成 Swagger 文档的 paths 部分，并按表分组"""
        paths = {}
        for table_name in self.data_engine.tables:
            # 为每个表创建一个分组
            self._swagger_tags.append({
                "name": table_name,
                "description": f"{table_name} 表的相关操作"
            })

            # 主路由：获取表数据（支持分页）
            paths[f'/api/{table_name}'] = {
                'get': {
                    'tags': [table_name],
                    'summary': f'获取 {table_name} 表数据',
                    'parameters': [
                        {
                            'name': 'page',
                            'in': 'query',
                            'required': False,
                            'schema': {
                                'type': 'integer',
                                'default': 1
                            },
                            'description': '页码，默认为 1'
                        },
                        {
                            'name': 'size',
                            'in': 'query',
                            'required': False,
                            'schema': {
                                'type': 'integer',
                                'default': 10
                            },
                            'description': '每页数据量，默认为 10'
                        }
                    ],
                    'responses': {
                        '200': {
                            'description': '成功返回表数据',
                            'content': {
                                'application/json': {
                                    'schema': {
                                        'type': 'array',
                                        'items': {
                                            'type': 'object',
                                            'properties': {
                                                'id': {'type': 'integer'},
                                                'name': {'type': 'string'},
                                                'age': {'type': 'integer'},
                                                'city': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            # 增删改查路由
            if self.enable_crud:
                # 添加数据
                paths[f'/api/{table_name}/add'] = {
                    'post': {
                        'tags': [table_name],
                        'summary': f'添加 {table_name} 表数据',
                        'requestBody': {
                            'content': {
                                'application/json': {
                                    'schema': {
                                        'type': 'object',
                                        'properties': {
                                            'id': {'type': 'integer'},
                                            'name': {'type': 'string'},
                                            'age': {'type': 'integer'},
                                            'city': {'type': 'string'}
                                        }
                                    }
                                }
                            }
                        },
                        'responses': {
                            '200': {
                                'description': '数据添加成功',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'message': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                # 更新数据（通过 ID）
                paths[f'/api/{table_name}/edit/{{id}}'] = {
                    'put': {
                        'tags': [table_name],
                        'summary': f'更新 {table_name} 表数据',
                        'parameters': [
                            {
                                'name': 'id',
                                'in': 'path',
                                'required': True,
                                'schema': {
                                    'type': 'integer'
                                }
                            }
                        ],
                        'requestBody': {
                            'content': {
                                'application/json': {
                                    'schema': {
                                        'type': 'object',
                                        'properties': {
                                            'id': {'type': 'integer'},
                                            'name': {'type': 'string'},
                                            'age': {'type': 'integer'},
                                            'city': {'type': 'string'}
                                        }
                                    }
                                }
                            }
                        },
                        'responses': {
                            '200': {
                                'description': '数据更新成功',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'message': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                # 删除数据（通过 ID）
                paths[f'/api/{table_name}/delete/{{id}}'] = {
                    'delete': {
                        'tags': [table_name],
                        'summary': f'删除 {table_name} 表数据',
                        'parameters': [
                            {
                                'name': 'id',
                                'in': 'path',
                                'required': True,
                                'schema': {
                                    'type': 'integer'
                                }
                            }
                        ],
                        'responses': {
                            '200': {
                                'description': '数据删除成功',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'message': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                # 获取单条数据（通过 ID）
                paths[f'/api/{table_name}/get/{{id}}'] = {
                    'get': {
                        'tags': [table_name],
                        'summary': f'获取 {table_name} 表单条数据',
                        'parameters': [
                            {
                                'name': 'id',
                                'in': 'path',
                                'required': True,
                                'schema': {
                                    'type': 'integer'
                                }
                            }
                        ],
                        'responses': {
                            '200': {
                                'description': '成功返回单条数据',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'id': {'type': 'integer'},
                                                'name': {'type': 'string'},
                                                'age': {'type': 'integer'},
                                                'city': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            },
                            '404': {
                                'description': '数据未找到',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'error': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                # 新增：通过字典条件删除数据
                paths[f'/api/{table_name}/delete'] = {
                    'post': {
                        'tags': [table_name],
                        'summary': f'删除 {table_name} 表数据（通过字典条件）',
                        'requestBody': {
                            'content': {
                                'application/json': {
                                    'schema': {
                                        'type': 'object',
                                        'properties': {
                                            'query': {
                                                'type': 'object',
                                                'description': '删除条件字典'
                                            }
                                        }
                                    }
                                }
                            }
                        },
                        'responses': {
                            '200': {
                                'description': '数据删除成功',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'message': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                # 新增：通过字典条件更新数据
                paths[f'/api/{table_name}/update'] = {
                    'post': {
                        'tags': [table_name],
                        'summary': f'更新 {table_name} 表数据（通过字典条件）',
                        'requestBody': {
                            'content': {
                                'application/json': {
                                    'schema': {
                                        'type': 'object',
                                        'properties': {
                                            'query': {
                                                'type': 'object',
                                                'description': '查询条件字典'
                                            },
                                            'update_data': {
                                                'type': 'object',
                                                'description': '更新数据字典'
                                            }
                                        }
                                    }
                                }
                            }
                        },
                        'responses': {
                            '200': {
                                'description': '数据更新成功',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'object',
                                            'properties': {
                                                'message': {'type': 'string'}
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                # 新增：通过字典条件查询数据
                paths[f'/api/{table_name}/find'] = {
                    'post': {
                        'tags': [table_name],
                        'summary': f'查询 {table_name} 表数据（通过字典条件）',
                        'requestBody': {
                            'content': {
                                'application/json': {
                                    'schema': {
                                        'type': 'object',
                                        'properties': {
                                            'query': {
                                                'type': 'object',
                                                'description': '查询条件字典'
                                            }
                                        }
                                    }
                                }
                            }
                        },
                        'responses': {
                            '200': {
                                'description': '成功返回查询结果',
                                'content': {
                                    'application/json': {
                                        'schema': {
                                            'type': 'array',
                                            'items': {
                                                'type': 'object',
                                                'properties': {
                                                    'id': {'type': 'integer'},
                                                    'name': {'type': 'string'},
                                                    'age': {'type': 'integer'},
                                                    'city': {'type': 'string'}
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

        # 将生成的路由信息合并到全局 Swagger 路径中
        self._swagger_paths.update(paths)

    def register_api_endpoints(self):
        """注册所有表的 API 路由"""
        # 生成默认的 Swagger 路径
        self.generate_swagger_paths()

        for table_name in self.data_engine.tables:
            self.register_table_api(table_name)

    def register_table_api(self, table_name: str):
        """为指定表注册 API 路由，并更新 Swagger 文档"""

        # 注册主路由，返回表数据（支持分页）
        @self.app.route(f'/api/{table_name}', methods=['GET'])
        def get_table_data():
            page = request.args.get('page', default=1, type=int)
            size = request.args.get('size', default=10, type=int)
            data = self.data_engine.findPage({}, page, size, table_name)
            return jsonify(data)

        # 如果启用 CRUD，注册增删改查接口
        if self.enable_crud:
            # 添加数据
            @self.app.route(f'/api/{table_name}/add', methods=['POST'])
            def add_handler():
                data = request.get_json()
                self.data_engine.save(data, table_name)
                return jsonify({"message": "数据添加成功"})

            # 更新数据（通过 ID）
            @self.app.route(f'/api/{table_name}/edit/<int:id>', methods=['PUT'])
            def edit_handler(id: int):
                data = request.get_json()
                data['id'] = id
                self.data_engine.save(data, table_name)
                return jsonify({"message": "数据更新成功"})

            # 删除数据（通过 ID）
            @self.app.route(f'/api/{table_name}/delete/<int:id>', methods=['DELETE'])
            def delete_handler(id: int):
                self.data_engine.deleteById(id, table_name)
                return jsonify({"message": "数据删除成功"})

            # 获取单条数据（通过 ID）
            @self.app.route(f'/api/{table_name}/get/<int:id>', methods=['GET'])
            def get_handler(id: int):
                data = self.data_engine.findById(id, table_name)
                if data:
                    return jsonify(data)
                else:
                    return jsonify({"error": "数据未找到"}), 404

            # 新增：通过字典条件删除数据
            @self.app.route(f'/api/{table_name}/delete', methods=['POST'])
            def delete_by_query_handler():
                query = request.get_json()
                if not query or not isinstance(query, dict):
                    return jsonify({"error": "请求体必须是一个非空字典"}), 400
                deleted_count = self.data_engine.delete(query, table_name)
                return jsonify({"message": f"删除了 {deleted_count} 条数据"})

            # 新增：通过字典条件更新数据
            @self.app.route(f'/api/{table_name}/update', methods=['POST'])
            def update_by_query_handler():
                request_data = request.get_json()
                if not request_data or not isinstance(request_data, dict):
                    return jsonify({"error": "请求体必须是一个非空字典"}), 400
                query = request_data.get("query")
                update_data = request_data.get("update_data")
                if not query or not update_data:
                    return jsonify({"error": "请求体必须包含 'query' 和 'update_data' 字段"}), 400
                updated_count = self.data_engine.update(query, update_data, table_name)
                return jsonify({"message": f"更新了 {updated_count} 条数据"})

            # 新增：通过字典条件查询数据
            @self.app.route(f'/api/{table_name}/find', methods=['POST'])
            def find_by_query_handler():
                query = request.get_json()
                if not query or not isinstance(query, dict):
                    return jsonify({"error": "请求体必须是一个非空字典"}), 400
                data = self.data_engine.find(query, table_name)
                return jsonify(data)

        # 如果启用文档，更新 Swagger 文档
        if self.enable_docs:
            self.generate_swagger_paths()

    def register_default_crud_routes(self, table_name: str):
        """注册默认的增删改查接口，支持字典操作"""

        @self.app.route(f'/api/{table_name}/add', methods=['POST'])
        def add_handler():
            data = request.get_json()
            self.data_engine.save(data, table_name)
            return jsonify({"message": "数据添加成功"})

        @self.app.route(f'/api/{table_name}/edit/<int:id>', methods=['PUT'])
        def edit_handler(id: int):
            data = request.get_json()
            data['id'] = id
            self.data_engine.save(data, table_name)
            return jsonify({"message": "数据更新成功"})

        @self.app.route(f'/api/{table_name}/delete/<int:id>', methods=['DELETE'])
        def delete_handler(id: int):
            self.data_engine.deleteById(id, table_name)
            return jsonify({"message": "数据删除成功"})

        @self.app.route(f'/api/{table_name}/get/<int:id>', methods=['GET'])
        def get_handler(id: int):
            data = self.data_engine.findById(id, table_name)
            if data:
                return jsonify(data)
            else:
                return jsonify({"error": "数据未找到"}), 404

        # 新增：通过字典条件删除数据
        @self.app.route(f'/api/{table_name}/delete', methods=['POST'])
        def delete_by_query_handler():
            query = request.get_json()
            if not query or not isinstance(query, dict):
                return jsonify({"error": "请求体必须是一个非空字典"}), 400
            deleted_count = self.data_engine.delete(query, table_name)
            return jsonify({"message": f"删除了 {deleted_count} 条数据"})

        # 新增：通过字典条件更新数据
        @self.app.route(f'/api/{table_name}/update', methods=['POST'])
        def update_by_query_handler():
            request_data = request.get_json()
            if not request_data or not isinstance(request_data, dict):
                return jsonify({"error": "请求体必须是一个非空字典"}), 400
            query = request_data.get("query")
            update_data = request_data.get("update_data")
            if not query or not update_data:
                return jsonify({"error": "请求体必须包含 'query' 和 'update_data' 字段"}), 400
            updated_count = self.data_engine.update(query, update_data, table_name)
            return jsonify({"message": f"更新了 {updated_count} 条数据"})

        # 新增：通过字典条件查询数据
        @self.app.route(f'/api/{table_name}/find', methods=['POST'])
        def find_by_query_handler():
            query = request.get_json()
            if not query or not isinstance(query, dict):
                return jsonify({"error": "请求体必须是一个非空字典"}), 400
            data = self.data_engine.find(query, table_name)
            return jsonify(data)

    def register_custom_crud_routes(self, table_name: str, crud):
        """使用自定义的 CRUD 类注册增删改查接口，支持字典操作"""

        # 添加数据
        @self.app.route(f'/api/{table_name}/add', methods=['POST'])
        def add_handler():
            data = request.get_json()
            return crud.add(data)

        # 更新数据（通过 ID）
        @self.app.route(f'/api/{table_name}/edit/<int:id>', methods=['PUT'])
        def edit_handler(id: int):
            data = request.get_json()
            return crud.edit(id, data)

        # 删除数据（通过 ID）
        @self.app.route(f'/api/{table_name}/delete/<int:id>', methods=['DELETE'])
        def delete_handler(id: int):
            return crud.delete(id)

        # 获取单条数据（通过 ID）
        @self.app.route(f'/api/{table_name}/get/<int:id>', methods=['GET'])
        def get_handler(id: int):
            return crud.get(id)

        # 新增：通过字典条件删除数据
        @self.app.route(f'/api/{table_name}/delete', methods=['POST'])
        def delete_by_query_handler():
            query = request.get_json()
            if not query or not isinstance(query, dict):
                return jsonify({"error": "请求体必须是一个非空字典"}), 400
            return crud.delete_by_query(query)

        # 新增：通过字典条件更新数据
        @self.app.route(f'/api/{table_name}/update', methods=['POST'])
        def update_by_query_handler():
            request_data = request.get_json()
            if not request_data or not isinstance(request_data, dict):
                return jsonify({"error": "请求体必须是一个非空字典"}), 400
            query = request_data.get("query")
            update_data = request_data.get("update_data")
            if not query or not update_data:
                return jsonify({"error": "请求体必须包含 'query' 和 'update_data' 字段"}), 400
            return crud.update_by_query(query, update_data)

        # 新增：通过字典条件查询数据
        @self.app.route(f'/api/{table_name}/find', methods=['POST'])
        def find_by_query_handler():
            query = request.get_json()
            if not query or not isinstance(query, dict):
                return jsonify({"error": "请求体必须是一个非空字典"}), 400
            return crud.find_by_query(query)

        # 如果启用文档，生成 Swagger 路径
        if self.enable_docs:
            self.generate_swagger_paths()

    def run_server(self, host: str = 'localhost', port: int = 5000, debug: bool = False):
        """启动 Flask 服务器"""
        self.app.run(host=host, port=port, debug=debug)


if __name__ == '__main__':
    # NebulaAPI 是一个基于 Flask 的轻量级 API 框架，旨在快速构建和管理数据驱动的 RESTful API。它的主要用途包括：
    #
    # 数据管理：
    #
    # 提供对数据表的增删改查（CRUD）操作。
    #
    # 支持通过字典条件进行灵活的查询、更新和删除操作。
    #
    # 自动化路由注册：
    #
    # 自动为数据表生成 RESTful API 路由。
    #
    # 支持手动注册路由，灵活控制 API 的行为。
    #
    # Swagger 文档支持：
    #
    # 自动生成 Swagger 文档，方便 API 使用者查看和测试。
    #
    # 灵活性：
    #
    # 支持多种数据源（如 Pandas DataFrame、文件路径、MongoDB 等）。
    #
    # 允许自定义 CRUD 逻辑，扩展功能。

    data = {
        'user': pd.DataFrame({
            'id': [1, 2, 3, 4],
            'name': ['Alice', 'Bob', 'Charlie', 'David'],
            'age': [24, 27, 22, 32],
            'city': ['New York', 'Los Angeles', 'Chicago', 'Houston']
        })
    }



    # 初始化 NebulaAPI（禁用自动注册路由）
    api = NebulaAPI(data_source=data, enable_docs=True, enable_crud=True, auto_register=False)

    # 手动注册路由
    api.register_table_api("user")

    # 启动服务器
    api.run_server(host='localhost', port=5000)